#!/bin/env python

from random import randint
import cv2
import numpy as np

windowName = "ANN demo"

# 创建人工神经网络
ann = cv2.ml.ANN_MLP_create()
# 设置训练方法
ann.setTrainMethod(cv2.ml.ANN_MLP_RPROP | cv2.ml.ANN_MLP_UPDATE_WEIGHTS)
# 设置激活函数
ann.setActivationFunction(cv2.ml.ANN_MLP_SIGMOID_SYM)
# 设置终止条件
ann.setTermCriteria((cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1))
# 创建拓扑结构
ann.setLayerSizes(np.array([3, 8, 4]))

"""Input arrays
weight, length, teeth
"""

"""Output arrays
dog, eagle, dolphin and dragon
"""
#狗，样本
def dog_sample():
  return [randint(5, 20), 1, randint(38, 42)]

#狗，标签
def dog_class():
  return [1, 0, 0, 0]

#秃鹰，样本
def condor_sample():
  return [randint(3,13), 3, 0]

#秃鹰，标签
def condor_class():
  return [0, 1, 0, 0]

#海豚，样本
def dolphin_sample():
  return [randint(30, 190), randint(5, 15), randint(80, 100)]

#海豚，标签
def dolphin_class():
  return [0, 0, 1, 0]

#龙，样本
def dragon_sample():
  return [randint(1200, 1800), randint(15, 40), randint(110, 180)]

#龙，标签
def dragon_class():
  return [0, 0, 0, 1]

#记录
def record(sample, classification):
  return (np.array([sample], dtype=np.float32), np.array([classification], dtype=np.float32))

records = []

'''
# 使用5000个样本训练
SAMPLES = 5000
for x in range(0, SAMPLES):
#   print ("Samples %d/%d" % (x, SAMPLES))
  ann.train(np.array([dog_sample()], dtype=np.float32), cv2.ml.ROW_SAMPLE, np.array([dog_class()], dtype=np.float32))
  ann.train(np.array([condor_sample()], dtype=np.float32), cv2.ml.ROW_SAMPLE, np.array([condor_class()], dtype=np.float32))
  ann.train(np.array([dolphin_sample()], dtype=np.float32), cv2.ml.ROW_SAMPLE, np.array([dolphin_class()], dtype=np.float32))
  ann.train(np.array([dragon_sample()], dtype=np.float32), cv2.ml.ROW_SAMPLE, np.array([dragon_class()], dtype=np.float32))
'''

#生成训练集
RECORDS = 5000
for x in range(0, RECORDS):
  records.append(record(dog_sample(), dog_class()))
  records.append(record(condor_sample(), condor_class()))
  records.append(record(dolphin_sample(), dolphin_class()))
  records.append(record(dragon_sample(), dragon_class()))

#迭代多轮，训练
EPOCHS = 5
for e in range(0, EPOCHS):
  print("Epoch %d:" % e)
  for t, c in records:
    ann.train(t, cv2.ml.ROW_SAMPLE, c)

dog_result = 0
for x in range(0, 100):
  clas = int(ann.predict(np.array([dog_sample()], dtype=np.float32))[0])
  if clas == 0:
    dog_result += 1

condor_result = 0
for x in range(0, 100):
  clas = int(ann.predict(np.array([condor_sample()], dtype=np.float32))[0])
  if clas == 1:
    condor_result += 1


dolphin_result = 0
for x in range(0, 100):
  clas = int(ann.predict(np.array([dolphin_sample()], dtype=np.float32))[0])
  if clas == 2:
    dolphin_result += 1

dragon_result = 0
for x in range(0, 100):
  clas = int(ann.predict(np.array([dragon_sample()], dtype=np.float32))[0])
  if clas == 3:
    dragon_result += 1

print ("Dog accuracy: %f%%" % (dog_result))
print ("condor accuracy: %f%%" % (condor_result))
print ("dolphin accuracy: %f%%" % (dolphin_result))
print ("dragon accuracy: %f%%" % (dragon_result))

if __name__ == '__main__':
    while True:
        break
        # if cv2.waitKey(0) == ord('q'):
        #     break

    # print('销毁窗口')
    # cv2.destroyAllWindows()
